/*
 *  BgSubtraction4.cpp
 *  opencvTest
 *
 *  Created by a1gucis on 3/1/09.
 *  Copyright 2009 __MyCompanyName__. All rights reserved.
 *
 */

#include "codebookmodel.h"
// Updates the codebook entry with a new data point
// 
// p            Pointer to a YUV pixel 
// c            Codebook for this pixel 
// cbBounds     Learning bounds for codebook (Rule of thumb: 10) 
// numChannels  Number of color channels we’re learning 
// 
// NOTES: 
//      cvBounds must be of length equal to numChannels 
// 
// RETURN 
//   codebook index 
// 
CodeBookModel::CodeBookModel(int height, int width){
	cbBounds = new unsigned[3];
	cbBounds[0] = 10;
	cbBounds[1] = 10;
	cbBounds[2] = 10;
	
	minMod = new int[3];
	minMod[0] = 10;
	minMod[1] = 10;
	minMod[2] = 10;
	
	maxMod = new int[3];
	maxMod[0] = 10;
	maxMod[1] = 10;
	maxMod[2] = 10;
	channels = 3;
	cb.t = 0;
	cb.numEntries = 0;
	codeBooks = new vector<CodeBook>(height*width);
}

CodeBookModel::~CodeBookModel() {
	delete[] cbBounds;
	delete[] minMod;
	delete[] maxMod;
	delete codeBooks;
	
}

int CodeBookModel::update_codebook(uchar *p, CodeBook& c, unsigned* cbBounds, int numChannels) {
	int i;
	c.t++;
	unsigned int high[3],low[3]; 
	for(int n=0; n<numChannels; n++) 
	{ 
		high[n] = *(p+n)+*(cbBounds+n); 
		if(high[n] > 255)
			high[n] = 255; 
		low[n] = *(p+n)-*(cbBounds+n); 
		if(low[n] < 0) low[n] = 0; 
	} 
	int matchChannel; 
	// SEE IF THIS FITS AN EXISTING CODEWORD 
	// 
	for(i=0; i<c.numEntries; i++){ 
		matchChannel = 0; 
		for(int n=0; n<numChannels; n++){ 
			if((c.cb[i]->learnLow[n] <= *(p+n)) && (*(p+n) <= c.cb[i]->learnHigh[n])) //Found an entry for this channel 
            { 
				matchChannel++; 
            } 
		} 
		if(matchChannel == numChannels) //If an entry was found 
		{ 
			c.cb[i]->t_last_update = c.t; 
			//adjust this codeword for the first channel 
			for(int n=0; n<numChannels; n++){ 
				if(c.cb[i]->max[n] < *(p+n)) 
				{ 
					c.cb[i]->max[n] = *(p+n); 
				} 
				else if(c.cb[i]->min[n] > *(p+n)) 
				{ 
					c.cb[i]->min[n] = *(p+n); 
				} 
			} 
			break;	
		} 
	}
	// OVERHEAD TO TRACK POTENTIAL STALE ENTRIES 
	// 
	for(int s=0; s<c.numEntries; s++){ 
		// Track which codebook entries are going stale: 
		// 
		int negRun = c.t - c.cb[s]->t_last_update; 
		if(c.cb[s]->stale < negRun) 
			c.cb[s]->stale = negRun; 
	} 
	
	// ENTER A NEW CODEWORD IF NEEDED 
	// 
	static int cbs = 0;
	if(i == c.numEntries) //if no existing codeword found, make one 
	{ 
		CodeElement **foo = new CodeElement* [c.numEntries+1]; 
		for(int ii=0; ii<c.numEntries; ii++) { 
			foo[ii] = c.cb[ii]; 
		} 
		foo[c.numEntries] = new CodeElement(); 
		if(c.numEntries) 
			delete [] c.cb; 
		c.cb = foo; 
		for(int n=0; n<numChannels; n++) { 
			c.cb[c.numEntries]->learnHigh[n] = high[n]; 
			c.cb[c.numEntries]->learnLow[n] = low[n]; 
			c.cb[c.numEntries]->max[n] = *(p+n); 
			c.cb[c.numEntries]->min[n] = *(p+n); 
		}
		c.cb[c.numEntries]->t_last_update = c.t; 
		c.cb[c.numEntries]->stale = 0; 
		c.numEntries += 1; 
	} 
	// SLOWLY ADJUST LEARNING BOUNDS 
	// 
	for(int n=0; n<numChannels; n++) 
	{ 
		if(c.cb[i]->learnHigh[n] < high[n]) c.cb[i]->learnHigh[n] += 1; 
		if(c.cb[i]->learnLow[n] > low[n]) c.cb[i]->learnLow[n] -= 1; 
	} 
	return i; 
}
// During learning, after you’ve learned for some period of time, 
// periodically call this to clear out stale codebook entries 
// 
// c   Codebook to clean up 
// 
// Return 
// number of entries cleared 
// 

int CodeBookModel::clear_stale_entries(CodeBook &c) {
	int staleThresh = c.t>>1; 
	int *keep = new int [c.numEntries]; 
	int keepCnt = 0; 
	// SEE WHICH CODEBOOK ENTRIES ARE TOO STALE 
	// 
	for(int i=0; i<c.numEntries; i++){ 
		if(c.cb[i]->stale > staleThresh) 
			keep[i] = 0; //Mark for destruction 
		else 
		{ 
			keep[i] = 1; //Mark to keep 
			keepCnt += 1;
		} 
	} 
	// KEEP ONLY THE GOOD 
	// 
	c.t = 0;         //Full reset on stale tracking 
	CodeElement **foo = new CodeElement* [keepCnt]; 
	int k=0; 
	for(int ii=0; ii<c.numEntries; ii++){ 
		if(keep[ii]) 
		{ 
			foo[k] = c.cb[ii]; 
			//We have to refresh these entries for next clearStale 
			foo[k]->t_last_update = 0; 
			k++; 
		} 
	} 
	// CLEAN UP 
	// 
	delete [] keep; 
	delete [] c.cb; 
	c.cb = foo; 
	int numCleared = c.numEntries - keepCnt; 
	c.numEntries = keepCnt; 
	c.t = 0;
	return(numCleared); 
	
}
// uchar background_diff( uchar *p, codeBook &c, 
//                         int minMod, int maxMod) 
// Given a pixel and a codebook, determine if the pixel is 
// covered by the codebook 
// 
// p            Pixel pointer (YUV interleaved) 
// c            Codebook reference 
// numChannels  Number of channels we are testing 
// maxMod       Add this (possibly negative) number onto
//              max level when determining if new pixel is foreground 
// minMod       Subract this (possibly negative) number from 
//              min level when determining if new pixel is foreground 
// 
// NOTES: 
// minMod and maxMod must have length numChannels, 
// e.g. 3 channels => minMod[3], maxMod[3]. There is one min and 
//      one max threshold per channel. 
// 
// Return 
// 0 => background, 255 => foreground 
// 

uchar CodeBookModel::background_diff(uchar *p, CodeBook &c, int numChannels, int* minMod, int* maxMod) {
	int matchChannel; 
	// SEE IF THIS FITS AN EXISTING CODEWORD 
	// 
	int i;
	for(i=0; i<c.numEntries; i++) { 
		matchChannel = 0; 
		for(int n=0; n<numChannels; n++) { 
			if((c.cb[i]->min[n] - minMod[n] <= *(p+n)) && 
			   (*(p+n) <= c.cb[i]->max[n] + maxMod[n])) { 
				matchChannel++; //Found an entry for this channel 
			} else { 
				break; 
			} 
		} 
		if(matchChannel == numChannels) { 
			break; //Found an entry that matched all channels 
		} 
	} 
	if(i >= c.numEntries) return(255); 
	return(0); 	
}

void CodeBookModel::setCbBounds(int x, int y, int z) {
	cbBounds[0] = x;
	cbBounds[1] = y;
	cbBounds[2] = z;
}

void CodeBookModel::updateBG(IplImage *image) {
	uchar *ptr;
	for (int i=0;i<image->height;i++) {
		for (int j=0;j<image->width;j++) {
			ptr = (uchar *)(image->imageData + i*image->widthStep+j*channels);
			update_codebook(ptr, (*codeBooks)[i*j], cbBounds, channels);
		}
	}
}

void CodeBookModel::clearStale() {
	for (int i=0;i<codeBooks->size();i++)
		clear_stale_entries((*codeBooks)[i]);
}

void CodeBookModel::bgDiff(IplImage *bg, IplImage *fg) {
	uchar *ptrBG, *ptrFG, pixelVal;
	for (int i=0;i<bg->height;i++) {
		for (int j=0;j<bg->width;j++) {
			ptrBG = (uchar *)(bg->imageData + i*bg->widthStep+j*channels);
			ptrFG = (uchar *)(fg->imageData + i*fg->widthStep+j);
			pixelVal = background_diff(ptrBG, (*codeBooks)[i*j], channels, minMod, maxMod);
			*ptrFG = pixelVal;
		}
	}
}